"""
Author: Morphlng
Date: 2024-03-14 17:00:18
LastEditTime: 2024-04-17 13:13:36
LastEditors: Morphlng
Description: Map interface for OpenDrive maps. Note that this module utilizes carla.Map API as the backend.
FilePath: /DrivingGym/src/driving_gym/simulation/common/hd_map.py
"""

from __future__ import annotations

import logging
import os
from collections import deque
from enum import Enum

import numpy as np
from shapely.geometry import Point, Polygon

from driving_gym.simulation.common.geom import Location, Rotation, Transform
from driving_gym.simulation.common.stub import try_import_carla

carla = try_import_carla()
logger = logging.getLogger(__name__)


class LaneMarkingColor(Enum):
    """Enum that defines the lane marking colors according to OpenDrive 1.4.

    The goal of this enum is to make sure that lane colors are correctly
    propogated from the simulator to env.

    https://github.com/carla-simulator/carla/blob/0.9.13/LibCarla/source/carla/road/element/LaneMarking.h#L39-L47
    """

    WHITE = 0
    BLUE = 1
    GREEN = 2
    RED = 3
    YELLOW = 4
    OTHER = 5


class LaneMarkingType(Enum):
    """Enum that defines the lane marking types according to OpenDrive 1.4.

    The goal of this enum is to make sure that lane markings are correctly
    propogated from the simulator to env.

    https://github.com/carla-simulator/carla/blob/0.9.13/LibCarla/source/carla/road/element/LaneMarking.h#L20-L37
    """

    OTHER = 0
    BROKEN = 1
    SOLID = 2
    SOLIDSOLID = 3
    SOLIDBROKEN = 4
    BROKENSOLID = 5
    BROKENBROKEN = 6
    BOTTSDOTS = 7
    GRASS = 8
    CURB = 9
    NONE = 10


class LaneChange(Enum):
    """Enum that defines the permission to turn either left, right, both or
    none for a given lane.

    The goal of this enum is to make sure that the lane change types are
    correctly propogated from the simulator to env.

    https://github.com/carla-simulator/carla/blob/0.9.13/LibCarla/source/carla/road/element/LaneMarking.h#L50-L55
    """

    NONE = 0
    RIGHT = 1
    LEFT = 2
    BOTH = 3


class LaneType(Enum):
    """Enum that defines the type of the lane according to OpenDrive 1.4.

    The goal of this enum is to make sure that the lane change types are
    correctly propogated from the simulator to env.

    https://github.com/carla-simulator/carla/blob/0.9.13/LibCarla/source/carla/road/Lane.h#L29-L52
    """

    NONE = 1
    DRIVING = 2
    STOP = 4
    SHOULDER = 8
    BIKING = 16
    SIDEWALK = 32
    BORDER = 64
    RESTRICTED = 128
    PARKING = 256
    BIDIRECTIONAL = 512
    MEDIAN = 1024
    SPECIAL1 = 2048
    SPECIAL2 = 4096
    SPECIAL3 = 8192
    ROADWORKS = 16384
    TRAM = 32768
    RAIL = 65536
    ENTRY = 131072
    EXIT = 262144
    OFFRAMP = 524288
    ONRAMP = 1048576
    ANY = -2  # 0xFFFFFFFE


class RoadOption(Enum):
    """Enum that defines the possible high-level route plans.

    RoadOptions are usually attached to waypoints we receive from
    the challenge environment.

    https://github.com/carla-simulator/carla/blob/0.9.13/PythonAPI/carla/agents/navigation/local_planner.py#L17-L28
    """

    VOID = -1
    LEFT = 1
    RIGHT = 2
    STRAIGHT = 3
    LANE_FOLLOW = 4
    CHANGE_LANE_LEFT = 5
    CHANGE_LANE_RIGHT = 6

    def __str__(self):
        return self.name


class LaneMarking(object):
    """Used to represent a lane marking.

    Attributes:
        marking_color (:py:class:`.LaneMarkingColor`): The color of the lane
            marking
        marking_type (:py:class:`.LaneMarkingType`): The type of the lane
            marking.
        lane_change (:py:class:`.LaneChange`): The type that defines the
            permission to either turn left, right, both or none.
    """

    def __init__(self, marking_color, marking_type, lane_change):
        self.marking_color = LaneMarkingColor(marking_color)
        self.marking_type = LaneMarkingType(marking_type)
        self.lane_change = LaneChange(lane_change)

    @classmethod
    def from_carla_lane_marking(cls, lane_marking):
        """Creates a LaneMarking from a carla lane marking.

        Args:
            lane_marking: An instance of a carla lane marking.

        Returns:
            :py:class:`.LaneMarking`: A lane-marking.
        """
        if not isinstance(lane_marking, carla.LaneMarking):
            raise ValueError(
                f"lane_marking should be of type `carla.LaneMarking`, getting {type(lane_marking)} instead"
            )

        return cls(lane_marking.color, lane_marking.type, lane_marking.lane_change)

    def __repr__(self):
        return "LaneMarking(color: {}, type: {}, change: {})".format(
            self.marking_color, self.marking_type, self.lane_change
        )


class Waypoint(object):
    """A waypoint in the world."""

    def __init__(
        self,
        id: int = -1,
        transform: Transform = None,
        road_id: int = -1,
        lane_id: int = -1,
        s: float = -1.0,
        section_id: int = -1,
        is_junction: bool = False,
        lane_width: float = 0,
    ):
        self.id: int = id
        self.transform: Transform = transform
        self.road_id: int = road_id
        self.lane_id: int = lane_id
        self.s: float = s
        self.section_id: int = section_id
        self.is_junction: bool = is_junction
        self.lane_width: float = lane_width
        self.lane_change: LaneChange = LaneChange.NONE
        self.lane_type: LaneType = LaneType.NONE
        self.left_lane_marking: LaneMarkingType = LaneMarkingType.NONE
        self.right_lane_marking: LaneMarkingType = LaneMarkingType.NONE

    @classmethod
    def from_carla_waypoint(cls, waypoint):
        """Creates a waypoint from a carla waypoint.

        Args:
            waypoint: carla waypoint.

        Returns:
            :py:class:`.Waypoint`: Waypoint object.
        """
        if not isinstance(waypoint, carla.Waypoint):
            raise ValueError(
                f"waypoint should be of type `carla.Waypoint`, getting {type(waypoint)} instead"
            )

        new_waypoint = cls()
        new_waypoint.id = waypoint.id
        new_waypoint.transform = Transform(
            Location(
                waypoint.transform.location.x,
                waypoint.transform.location.y,
                waypoint.transform.location.z,
            ),
            Rotation(
                waypoint.transform.rotation.pitch,
                waypoint.transform.rotation.yaw,
                waypoint.transform.rotation.roll,
            ),
        )
        new_waypoint.road_id = waypoint.road_id
        new_waypoint.lane_id = waypoint.lane_id
        new_waypoint.s = waypoint.s
        new_waypoint.section_id = waypoint.section_id
        new_waypoint.is_junction = waypoint.is_junction
        new_waypoint.lane_width = waypoint.lane_width
        new_waypoint.lane_change = LaneChange(waypoint.lane_change)
        new_waypoint.lane_type = LaneType(waypoint.lane_type)
        new_waypoint.left_lane_marking = LaneMarking.from_carla_lane_marking(
            waypoint.left_lane_marking
        )
        new_waypoint.right_lane_marking = LaneMarking.from_carla_lane_marking(
            waypoint.right_lane_marking
        )
        return new_waypoint

    def __repr__(self):
        return "Waypoint(transform={}, road_id={}, lane_id={}, s={})".format(
            self.transform, self.road_id, self.lane_id, self.s
        )


class Lane(object):
    """Stores information about a lane.

    Args:
        id (:obj:`int`): The id of the lane (0 for ego lane, negative for
            left lanes, and positive for right lanes).
        left_markings: List of lane mark locations.
        right_markings: List of lane mark locations.
    """

    def __init__(
        self, id: int, left_markings: list[Location], right_markings: list[Location]
    ):
        self.id = id
        self.left_markings = left_markings
        self.right_markings = right_markings
        self._lane_polygon = None

    def get_closest_lane_waypoint(self, location: Location):
        if self.is_on_lane(location):
            return Transform(location, Rotation())
        closest_location = None
        min_dist = np.infty
        for loc in self.left_markings:
            dist = loc.distance(location)
            if dist < min_dist:
                min_dist = dist
                closest_location = loc
        for loc in self.right_markings:
            dist = loc.distance(location)
            if dist < min_dist:
                min_dist = dist
                closest_location = loc
        return Transform(closest_location, Rotation())

    def get_lane_center_transforms(self):
        if len(self.left_markings) < len(self.right_markings):
            anchor_markings = self.left_markings
            other_markings = self.right_markings
        else:
            anchor_markings = self.right_markings
            other_markings = self.left_markings
        index_other = 0
        center_markings = deque([])
        for loc in anchor_markings:
            dist = loc.distance(other_markings[index_other])
            while index_other + 1 < len(other_markings) and dist > loc.distance(
                other_markings[index_other + 1]
            ):
                index_other += 1
                dist = loc.distance(other_markings[index_other])
            if index_other < len(other_markings):
                other_loc = other_markings[index_other]
                center_location = Location(
                    (loc.x + other_loc.x) / 2.0,
                    (loc.y + other_loc.y) / 2.0,
                    (loc.z + other_loc.z) / 2.0,
                )
                center_markings.append(Transform(center_location, Rotation()))
        return center_markings

    def is_on_lane(self, location: Location):
        # We only create the lane polygon if it is necessary.
        if not self._lane_polygon:
            self._create_lane_polygon()
        return self._lane_polygon.contains(Point(location.x, location.y))

    def _create_lane_polygon(self):
        points = [(0, self.left_markings[0].y)]
        for loc in self.left_markings:
            points.append((loc.x, loc.y))
        for loc in reversed(self.right_markings):
            points.append((loc.x, loc.y))
        points.append((0, self.right_markings[0].y))
        self._lane_polygon = Polygon(points)

    def __repr__(self):
        return "Lane(id: {}, {})".format(
            self.id, zip(self.left_markings, self.right_markings)
        )


class HDMap(object):
    """A map interface for OpenDrive maps. This module utilize carla.Map API as the backend."""

    def __init__(self, simulator_map):
        """Create a map from a simulator map (Either a carla.Map instance or str)

        Args:
            simulator_map (carla.Map | str): The simulator map to create the map from.
            If it is a string, it can either be a path to a OpenDrive file or the
            content of the OpenDrive file.
        """
        from driving_gym.simulation.common.global_route_planner import (
            GlobalRoutePlanner,
        )

        if isinstance(simulator_map, carla.Map):
            self._map = simulator_map
        elif os.path.isfile(simulator_map):
            try:
                with open(simulator_map, "r") as file:
                    self._map = carla.Map(
                        os.path.splitext(os.path.basename(simulator_map))[0],
                        file.read(),
                    )
            except Exception as e:
                raise ValueError(f"Failed to load map from file: {e}")
        elif isinstance(simulator_map, str):
            try:
                self._map = carla.Map("Test", simulator_map)
            except Exception as e:
                raise ValueError(f"Failed to load map from string: {e}")
        else:
            raise ValueError("Invalid map input")

        self._grp = GlobalRoutePlanner(self._map, 1.0)

    def get_closest_lane_waypoint(self, location: Location) -> Waypoint:
        """Returns the road closest waypoint to location.

        Args:
            location (:py:class:`geom.Location`): Location in world
                coordinates.

        Returns:
            :py:class:`.Waypoint`: Waypoint or None if no
            waypoint is found.
        """
        waypoint = self._get_waypoint(location, project_to_road=True)
        if waypoint is not None:
            return Waypoint.from_carla_waypoint(waypoint)
        else:
            return None

    def is_intersection(self, location: Location) -> bool:
        """Checks if a location is in an intersection.

        Args:
            location (:py:class:`geom.Location`): Location in world
                coordinates.

        Returns:
            bool: True if the location is in an intersection.
        """
        waypoint = self._get_waypoint(location)
        if waypoint is None:
            # The map didn't return a waypoint because the location not within
            # mapped location.
            return False
        else:
            return self.__is_intersection(waypoint)

    def __is_intersection(self, waypoint) -> bool:
        if waypoint.is_junction:
            return True
        if getattr(waypoint, "is_intersection", False):
            return True
        return False

    def is_on_lane(self, location: Location) -> bool:
        """Checks if a location is on a lane.

        Args:
            location (:py:class:`geom.Location`): Location in world coordinates.

        Returns:
            bool: True if the location is on a lane.
        """
        waypoint = self._get_waypoint(location, lane_type=LaneType.ANY)
        if waypoint is None or waypoint.lane_type != carla.LaneType.Driving:
            # The map didn't return a waypoint because the location not within
            # mapped location.
            return False
        else:
            return True

    def are_on_same_lane(self, location1: Location, location2: Location) -> bool:
        """Checks if two locations are on the same lane.

        Args:
            location1 (:py:class:`geom.Location`): Location in world
                coordinates.
            location2 (:py:class:`geom.Location`): Location in world
                coordinates.

        Returns:
            bool: True if the two locations are on the same lane.
        """
        waypoint1 = self._get_waypoint(location1, lane_type=LaneType.DRIVING)
        if waypoint1 is None:
            # First location is not on a drivable lane.
            return False
        waypoint2 = self._get_waypoint(location2, lane_type=LaneType.DRIVING)
        if waypoint2 is None:
            # Second location is not on a drivable lane.
            return False
        if waypoint1.road_id == waypoint2.road_id:
            return waypoint1.lane_id == waypoint2.lane_id
        else:
            # Return False if we're in intersection and the other
            # obstacle isn't.
            if self.__is_intersection(waypoint1) and not self.__is_intersection(
                waypoint2
            ):
                return False
            if waypoint2.lane_type == LaneType.DRIVING.value:
                # This may return True when the lane is different, but in
                # with a different road_id.
                # TODO: Figure out how lane id map across road id.
                return True
        return False

    def is_on_opposite_lane(self, transform: Transform):
        """Checks if a transform is on an opposite lane.

        Args:
            transform (:py:class:`geom.Transform`): Transform in world coordinates.

        Returns:
            bool: True if the transform is on the opposite lane.
        """
        waypoint = self._get_waypoint(transform.location, lane_type=LaneType.DRIVING)
        if waypoint is None:
            return True
        if self.__is_intersection(waypoint):
            return False

        # TODO: 90 degrees?
        if abs(waypoint.transform.rotation.yaw - transform.rotation.yaw) > 140:
            return True
        else:
            return False

    def is_at_stop(self, location: Location):
        """Checks if a location is close to a stop sign.

        Args:
            location (:py:class:`geom.Location`): Location in world
                coordinates.

        Returns:
            bool: True if the location is at a stop sign.
        """
        # TODO: OpenDrive 1.4 does bit have stop signs.
        waypoint = self._get_waypoint(location, lane_type=LaneType.STOP)
        return waypoint is not None

    def distance_to_intersection(
        self, location: Location, max_distance_to_check: float = 30
    ):
        """Computes the distance (in meters) from location to an intersection.

        The method starts from location, moves forward until it reaches an
        intersection or exceeds max_distance_to_check.

        Args:
            location (:py:class:`geom.Location`): The starting location
                in world coordinates.
            max_distance_to_check (:obj:`int`): Max distance to move forward (in meters).

        Returns:
            :obj:`int`: The distance in meters, or None if there is no
            intersection within max_distance_to_check.
        """
        waypoint = self._get_waypoint(location)
        if waypoint is None:
            return None
        # We're already in an intersection.
        if self.__is_intersection(waypoint):
            return 0
        for i in range(1, max_distance_to_check + 1):
            waypoints = waypoint.next(1)
            if waypoints is None or len(waypoints) == 0:
                return None
            for w in waypoints:
                if self.__is_intersection(w):
                    return i
            waypoint = waypoints[0]
        return None

    def is_on_bidirectional_lane(self, location: Location) -> bool:
        """Checks if a location is a bidirectional lane.

        Args:
            location (:py:class:`geom.Location`): Location in world coordinates.

        Returns:
            bool: True if the location is on a bidirectional lane.
        """
        waypoint = self._get_waypoint(location, lane_type=LaneType.BIDIRECTIONAL)
        return waypoint is not None

    def must_obey_traffic_light(
        self, ego_location: Location, tl_location: Location
    ) -> bool:
        """Checks if an ego vehicle must obey a traffic light.

        Args:
            ego_location (:py:class:`geom.Location`): Location of the
                ego vehicle in world coordinates.
            tl_location (:py:class:`geom.Location`): Location of the
                traffic light in world coordinates.

        Returns:
            bool: True if the ego vehicle must obey the traffic light.
        """
        waypoint = self._get_waypoint(ego_location)
        if waypoint is not None and self.__is_intersection(waypoint):
            # Do not obey traffic light if ego is already in the intersection.
            return False

        # TODO: Use the following rules.

        return True

    def _must_obey_european_traffic_light(
        self,
        ego_transform: Transform,
        tl_locations: list[Location],
        tl_max_dist_thresh: float,
    ) -> bool:
        ego_waypoint = self._get_waypoint(ego_transform.location)
        # We're not on a road, or we're already in the intersection. Carry on.
        if ego_waypoint is None or self.__is_intersection(ego_waypoint):
            return (False, None)
        # Iterate through traffic lights.
        for tl_loc in tl_locations:
            tl_waypoint = self._get_waypoint(tl_loc)
            if (
                tl_waypoint.road_id != ego_waypoint.road_id
                or tl_waypoint.lane_id != ego_waypoint.lane_id
            ):
                continue
            if ego_transform.is_within_distance_ahead(tl_loc, tl_max_dist_thresh):
                return (True, tl_loc)
        return (False, None)

    def _must_obey_american_traffic_light(
        self,
        ego_transform: Transform,
        tl_locations: list[Location],
        tl_max_dist_thresh: float,
    ) -> bool:
        ego_waypoint = self._get_waypoint(ego_transform.location)
        # We're not on a road, or we're already in the intersection. Carry on.
        if ego_waypoint is None or self.__is_intersection(ego_waypoint):
            return (False, None)

        min_angle = 25.0
        selected_tl_loc = None
        for tl_loc in tl_locations:
            if ego_transform.is_within_distance_ahead(tl_loc, tl_max_dist_thresh):
                angle, distance = ego_transform.get_angle_and_magnitude(tl_loc)
                if distance < 60.0 and angle < min(25.0, min_angle):
                    min_angle = angle
                    selected_tl_loc = tl_loc
        if selected_tl_loc is not None:
            return (True, selected_tl_loc)
        else:
            return (False, None)

    def get_lane(
        self, location: Location, waypoint_precision: float = 0.05, lane_id: int = 0
    ) -> Lane:
        lane_waypoints = []
        # Consider waypoints in opposite direction of camera so we can get
        # lane data for adjacent lanes in opposing directions.
        previous_wp = [self._get_waypoint(location)]

        while len(previous_wp) == 1:
            lane_waypoints.append(previous_wp[0])
            previous_wp = previous_wp[0].previous(waypoint_precision)

        next_wp = [self._get_waypoint(location)]

        while len(next_wp) == 1:
            lane_waypoints.append(next_wp[0])
            next_wp = next_wp[0].next(waypoint_precision)

        # Get the left and right markings of the lane and send it as a message.
        left_markings = [
            self.lateral_shift(w.transform, -w.lane_width * 0.5) for w in lane_waypoints
        ]
        right_markings = [
            self.lateral_shift(w.transform, w.lane_width * 0.5) for w in lane_waypoints
        ]
        return Lane(lane_id, left_markings, right_markings)

    def get_left_lane(self, location: Location) -> Waypoint:
        waypoint = self._get_waypoint(location)
        if waypoint is not None:
            left_lane_waypoint = waypoint.get_left_lane()
            if left_lane_waypoint:
                return Waypoint.from_carla_waypoint(left_lane_waypoint)
        return None

    def get_right_lane(self, location: Location) -> Waypoint:
        waypoint = self._get_waypoint(location)
        if waypoint is not None:
            right_lane_waypoint = waypoint.get_right_lane()
            if right_lane_waypoint:
                return Waypoint.from_carla_waypoint(right_lane_waypoint)
        return None

    def get_all_lanes(self, location: Location):
        lanes = [self.get_lane(location)]

        waypoint = self._get_waypoint(location)
        if waypoint is not None:
            wp_left = waypoint.get_left_lane()
            w_rotation = waypoint.transform.rotation
            while wp_left and wp_left.lane_type == LaneType.DRIVING.value:
                left_location = Location(
                    wp_left.transform.location.x,
                    wp_left.transform.location.y,
                    wp_left.transform.location.z,
                )
                lanes.append(self.get_lane(left_location, lane_id=wp_left.lane_id))

                # If left lane is facing the opposite direction, its left
                # lane would point back to the current lane, so we select
                # its right lane to get the left lane relative to current.
                if w_rotation == wp_left.transform.rotation:
                    wp_left = wp_left.get_left_lane()
                else:
                    wp_left = wp_left.get_right_lane()

            wp_right = waypoint.get_right_lane()
            while wp_right and wp_right.lane_type == LaneType.DRIVING.value:
                right_location = Location(
                    wp_right.transform.location.x,
                    wp_right.transform.location.y,
                    wp_right.transform.location.z,
                )
                lanes.append(self.get_lane(right_location, lane_id=wp_right.lane_id))

                # Same logic as above. If right lane of current is in
                # opposite direction, move rightwards by selecting it's
                # left lane.
                if w_rotation == wp_right.transform.rotation:
                    wp_right = wp_right.get_right_lane()
                else:
                    wp_right = wp_left.get_left_lane()
        return lanes

    def compute_waypoints(self, source_loc: Location, destination_loc: Location):
        """Computes waypoints between two locations.

        Assumes that the ego vehicle has the same orientation as the lane on
        which it is on.

        Args:
            source_loc (:py:class:`geom.Location`): Source location in
                world coordinates.
            destination_loc (:py:class:`geom.Location`): Destination
                location in world coordinates.

        Returns:
            list(:py:class:`.Waypoint`): List of waypoints.
        """
        start_waypoint = self._get_waypoint(
            source_loc, project_to_road=True, lane_type=LaneType.DRIVING
        )
        end_waypoint = self._get_waypoint(
            destination_loc, project_to_road=True, lane_type=LaneType.DRIVING
        )
        assert start_waypoint and end_waypoint, "Map could not find waypoints"
        route = self._grp.trace_route(
            start_waypoint.transform.location, end_waypoint.transform.location
        )
        # TODO: The planner returns several options in intersections.
        # We always take the first one, but this is not correct.
        return deque([Waypoint.from_carla_waypoint(waypoint[0]) for waypoint in route])

    def lateral_shift(self, transform: Transform, shift: float) -> Location:
        transform.rotation.yaw += 90
        shifted = transform.location + shift * transform.get_forward_vector()
        return Location(shifted.x, shifted.y, shifted.z)

    def _get_waypoint(
        self,
        location: Location,
        project_to_road: bool = True,
        lane_type: LaneType = LaneType.ANY,
    ):
        try:
            waypoint = self._map.get_waypoint(
                carla.Location(location.x, location.y, location.z),
                project_to_road=project_to_road,
                lane_type=lane_type.value,
            )
        except RuntimeError as err:
            logger.error(f"Failed to get waypoint: {err}")
            waypoint = None
        return waypoint
